{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/beta/ESMFold_advanced.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AMwhSNWEzRCZ"
      },
      "source": [
        "#**ESMFold_advanced**\n",
        "for more details see: [Github](https://github.com/facebookresearch/esm/tree/main/esm), [Preprint](https://www.biorxiv.org/content/10.1101/2022.07.20.500902v1)\n",
        "\n",
        "#### **Tips and Instructions**\n",
        "- click the little ▶ play icon to the left of each cell below.\n",
        "\n",
        "#### **Colab Limitations**\n",
        "- On Tesla T4 (typical free colab GPU), max total length ~ 900\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "catNptjYwNM6",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "%%time\n",
        "#@title install\n",
        "#@markdown install ESMFold, OpenFold and download Params (~2min 30s)\n",
        "\n",
        "import os, time\n",
        "if not os.path.isfile(\"esmfold.model\"):\n",
        "  # download esmfold params\n",
        "  os.system(\"apt-get install aria2 -qq\")\n",
        "  os.system(\"aria2c -q -x 16 https://colabfold.steineggerlab.workers.dev/esm/esmfold.model &\")\n",
        "\n",
        "  # install libs\n",
        "  os.system(\"pip install -q omegaconf pytorch_lightning biopython ml_collections einops py3Dmol\")\n",
        "  os.system(\"pip install -q git+https://github.com/NVIDIA/dllogger.git\")\n",
        "\n",
        "  # install openfold\n",
        "  commit = \"6908936b68ae89f67755240e2f588c09ec31d4c8\"\n",
        "  os.system(f\"pip install -q git+https://github.com/aqlaboratory/openfold.git@{commit}\")\n",
        "\n",
        "  # install esmfold\n",
        "  os.system(f\"pip install -q git+https://github.com/sokrypton/esm.git@beta\")\n",
        "\n",
        "  # wait for Params to finish downloading...\n",
        "  if not os.path.isfile(\"esmfold.model\"):\n",
        "    # backup source!\n",
        "    os.system(\"aria2c -q -x 16 https://files.ipd.uw.edu/pub/esmfold/esmfold.model\")\n",
        "  else:\n",
        "    while os.path.isfile(\"esmfold.model.aria2\"):\n",
        "      time.sleep(5)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2rqTrvifh17B",
        "cellView": "form"
      },
      "outputs": [],
      "source": [
        "#@title ##run **ESMFold**\n",
        "%%time\n",
        "from string import ascii_uppercase, ascii_lowercase\n",
        "import hashlib, re, os\n",
        "import numpy as np\n",
        "from jax.tree_util import tree_map\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.special import softmax\n",
        "\n",
        "def parse_output(output):\n",
        "  pae = (output[\"aligned_confidence_probs\"][0] * np.arange(64)).mean(-1) * 31\n",
        "  plddt = output[\"plddt\"][0,:,1]\n",
        "  \n",
        "  bins = np.append(0,np.linspace(2.3125,21.6875,63))\n",
        "  sm_contacts = softmax(output[\"distogram_logits\"],-1)[0]\n",
        "  sm_contacts = sm_contacts[...,bins<8].sum(-1)\n",
        "  xyz = output[\"positions\"][-1,0,:,1]\n",
        "  mask = output[\"atom37_atom_exists\"][0,:,1] == 1\n",
        "  o = {\"pae\":pae[mask,:][:,mask],\n",
        "       \"plddt\":plddt[mask],\n",
        "       \"sm_contacts\":sm_contacts[mask,:][:,mask],\n",
        "       \"xyz\":xyz[mask]}\n",
        "  if \"contacts\" in output[\"lm_output\"]:\n",
        "    lm_contacts = output[\"lm_output\"][\"contacts\"].astype(float)[0]\n",
        "    o[\"lm_contacts\"] = lm_contacts[mask,:][:,mask]\n",
        "  return o\n",
        "\n",
        "def get_hash(x): return hashlib.sha1(x.encode()).hexdigest()\n",
        "alphabet_list = list(ascii_uppercase+ascii_lowercase)\n",
        "\n",
        "jobname = \"test\" #@param {type:\"string\"}\n",
        "jobname = re.sub(r'\\W+', '', jobname)[:50]\n",
        "\n",
        "sequence = \"GWSTELEKHREELKEFLKKEGITNVEIRIDNGRLEVRVEGGTERLKRFLEELRQKLEKKGYTVDIKIE\" #@param {type:\"string\"}\n",
        "sequence = re.sub(\"[^A-Z:]\", \"\", sequence.replace(\"/\",\":\").upper())\n",
        "sequence = re.sub(\":+\",\":\",sequence)\n",
        "sequence = re.sub(\"^[:]+\",\"\",sequence)\n",
        "sequence = re.sub(\"[:]+$\",\"\",sequence)\n",
        "\n",
        "#@markdown ---\n",
        "#@markdown ###**Advanced Options**\n",
        "num_recycles = 3 #@param [\"0\", \"1\", \"2\", \"3\", \"6\", \"12\"] {type:\"raw\"}\n",
        "get_LM_contacts = False #@param {type:\"boolean\"}\n",
        "\n",
        "#@markdown **multimer options (experimental)**\n",
        "#@markdown - use \"/\" to specify chainbreaks, (eg. sequence=\"AAA/AAA\")\n",
        "#@markdown - for homo-oligomeric predictions, set copies > 1\n",
        "\n",
        "copies = 1 #@param {type:\"integer\"}\n",
        "chain_linker = 25 #@param {type:\"number\"}\n",
        "if copies == \"\" or copies <= 0: copies = 1\n",
        "sequence = \":\".join([sequence] * copies)\n",
        "\n",
        "#@markdown **sampling options (experimental)**\n",
        "#@markdown - Samples are generated via random masking (defined by `masking_rate`) \n",
        "#@markdown of input sequence (stochastic_mode=\"LM\") and/or via dropout within structure module (stochastic_mode=\"SM\").\n",
        "samples = None #@param [\"None\", \"1\", \"4\", \"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
        "masking_rate = 0.15 #@param {type:\"number\"}\n",
        "stochastic_mode = \"LM\" #@param [\"LM\", \"LM_SM\", \"SM\"]\n",
        "\n",
        "ID = jobname+\"_\"+get_hash(sequence)[:5]\n",
        "seqs = sequence.split(\":\")\n",
        "lengths = [len(s) for s in seqs]\n",
        "length = sum(lengths)\n",
        "print(\"length\",length)\n",
        "\n",
        "u_seqs = list(set(seqs))\n",
        "if len(seqs) == 1: mode = \"mono\"\n",
        "elif len(u_seqs) == 1: mode = \"homo\"\n",
        "else: mode = \"hetero\"\n",
        "\n",
        "if \"model\" not in dir():\n",
        "  import torch\n",
        "  model = torch.load(\"esmfold.model\")\n",
        "  model.cuda().requires_grad_(False)\n",
        "\n",
        "# optimized for Tesla T4\n",
        "if length > 700:\n",
        "  model.trunk.set_chunk_size(64)\n",
        "else:\n",
        "  model.trunk.set_chunk_size(128)\n",
        "\n",
        "best_pdb_str = None\n",
        "best_ptm = 0\n",
        "best_output = None\n",
        "traj = []\n",
        "\n",
        "num_samples = 1 if samples is None else samples\n",
        "for seed in range(num_samples):\n",
        "  torch.cuda.empty_cache()\n",
        "  if samples is None:\n",
        "    seed = \"default\"\n",
        "    mask_rate = 0.0\n",
        "    model.train(False)\n",
        "  else:\n",
        "    torch.manual_seed(seed)\n",
        "    mask_rate = masking_rate if \"LM\" in stochastic_mode else 0.0\n",
        "    model.train(\"SM\" in stochastic_mode)\n",
        "\n",
        "  output = model.infer(sequence,\n",
        "                       num_recycles=num_recycles,\n",
        "                       chain_linker=\"X\"*chain_linker,\n",
        "                       residue_index_offset=512,\n",
        "                       mask_rate=mask_rate,\n",
        "                       return_contacts=get_LM_contacts)\n",
        "  \n",
        "  pdb_str = model.output_to_pdb(output)[0]\n",
        "  output = tree_map(lambda x: x.cpu().numpy(), output)\n",
        "  ptm = output[\"ptm\"][0]\n",
        "  plddt = output[\"plddt\"][0,:,1].mean()\n",
        "  traj.append(parse_output(output))\n",
        "  print(f'{seed} ptm: {ptm:.3f} plddt: {plddt:.1f}')\n",
        "  if ptm > best_ptm:\n",
        "    best_pdb_str = pdb_str\n",
        "    best_ptm = ptm\n",
        "    best_output = output\n",
        "  os.system(f\"mkdir -p {ID}\")\n",
        "  if samples is None:\n",
        "    pdb_filename = f\"{ID}/ptm{ptm:.3f}_r{num_recycles}_{seed}.pdb\"\n",
        "  else:\n",
        "    pdb_filename = f\"{ID}/ptm{ptm:.3f}_r{num_recycles}_seed{seed}_{stochastic_mode}_m{masking_rate:.2f}.pdb\"\n",
        "\n",
        "  with open(pdb_filename,\"w\") as out:\n",
        "    out.write(pdb_str)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tg7eFMC-TYRX"
      },
      "outputs": [],
      "source": [
        "#@title display (optional) {run: \"auto\"}\n",
        "#@markdown Note: If samples selected, the model with max pTM is displayed.\n",
        "import py3Dmol\n",
        "\n",
        "\n",
        "pymol_color_list = [\"#33ff33\",\"#00ffff\",\"#ff33cc\",\"#ffff00\",\"#ff9999\",\"#e5e5e5\",\"#7f7fff\",\"#ff7f00\",\n",
        "                    \"#7fff7f\",\"#199999\",\"#ff007f\",\"#ffdd5e\",\"#8c3f99\",\"#b2b2b2\",\"#007fff\",\"#c4b200\",\n",
        "                    \"#8cb266\",\"#00bfbf\",\"#b27f7f\",\"#fcd1a5\",\"#ff7f7f\",\"#ffbfdd\",\"#7fffff\",\"#ffff7f\",\n",
        "                    \"#00ff7f\",\"#337fcc\",\"#d8337f\",\"#bfff3f\",\"#ff7fff\",\"#d8d8ff\",\"#3fffbf\",\"#b78c4c\",\n",
        "                    \"#339933\",\"#66b2b2\",\"#ba8c84\",\"#84bf00\",\"#b24c66\",\"#7f7f7f\",\"#3f3fa5\",\"#a5512b\"]\n",
        "\n",
        "def show_pdb(pdb_str, show_sidechains=False, show_mainchains=False,\n",
        "             color=\"pLDDT\", chains=None, vmin=50, vmax=90,\n",
        "             size=(800,480), hbondCutoff=4.0,\n",
        "             Ls=None,\n",
        "             animate=False):\n",
        "  \n",
        "  if chains is None:\n",
        "    chains = 1 if Ls is None else len(Ls)\n",
        "  view = py3Dmol.view(js='https://3dmol.org/build/3Dmol.js', width=size[0], height=size[1])\n",
        "  if animate:\n",
        "    view.addModelsAsFrames(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "  else:\n",
        "    view.addModel(pdb_str,'pdb',{'hbondCutoff':hbondCutoff})\n",
        "  if color == \"pLDDT\":\n",
        "    view.setStyle({'cartoon': {'colorscheme': {'prop':'b','gradient': 'roygb','min':vmin,'max':vmax}}})\n",
        "  elif color == \"rainbow\":\n",
        "    view.setStyle({'cartoon': {'color':'spectrum'}})\n",
        "  elif color == \"chain\":\n",
        "    for n,chain,color in zip(range(chains),alphabet_list,pymol_color_list):\n",
        "       view.setStyle({'chain':chain},{'cartoon': {'color':color}})\n",
        "  if show_sidechains:\n",
        "    BB = ['C','O','N']\n",
        "    view.addStyle({'and':[{'resn':[\"GLY\",\"PRO\"],'invert':True},{'atom':BB,'invert':True}]},\n",
        "                  {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"GLY\"},{'atom':'CA'}]},\n",
        "                  {'sphere':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "    view.addStyle({'and':[{'resn':\"PRO\"},{'atom':['C','O'],'invert':True}]},\n",
        "                  {'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})  \n",
        "  if show_mainchains:\n",
        "    BB = ['C','O','N','CA']\n",
        "    view.addStyle({'atom':BB},{'stick':{'colorscheme':f\"WhiteCarbon\",'radius':0.3}})\n",
        "  view.zoomTo()\n",
        "  if animate: view.animate()\n",
        "  return view\n",
        "\n",
        "color = \"confidence\" #@param [\"confidence\", \"rainbow\", \"chain\"]\n",
        "if color == \"confidence\": color = \"pLDDT\"\n",
        "show_sidechains = False #@param {type:\"boolean\"}\n",
        "show_mainchains = False #@param {type:\"boolean\"}\n",
        "show_pdb(best_pdb_str, color=color,\n",
        "         show_sidechains=show_sidechains,\n",
        "         show_mainchains=show_mainchains,\n",
        "         Ls=lengths).show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hkiKtX43voI-"
      },
      "outputs": [],
      "source": [
        "#@title plot confidence (optional)\n",
        "#@markdown Note: If samples selected, the model with max pTM is displayed.\n",
        "\n",
        "dpi = 100 #@param {type:\"integer\"}\n",
        "\n",
        "def plot_ticks(Ls):\n",
        "  Ln = sum(Ls)\n",
        "  L_prev = 0\n",
        "  for L_i in Ls[:-1]:\n",
        "    L = L_prev + L_i\n",
        "    L_prev += L_i\n",
        "    plt.plot([0,Ln],[L,L],color=\"black\")\n",
        "    plt.plot([L,L],[0,Ln],color=\"black\")\n",
        "  ticks = np.cumsum([0]+Ls)\n",
        "  ticks = (ticks[1:] + ticks[:-1])/2\n",
        "  plt.yticks(ticks,alphabet_list[:len(ticks)])\n",
        "\n",
        "def plot_confidence(output, Ls=None, dpi=100):\n",
        "  O = parse_output(output)\n",
        "  if \"lm_contacts\" in O:\n",
        "    plt.figure(figsize=(20,4), dpi=dpi)\n",
        "    plt.subplot(1,4,1)\n",
        "  else:\n",
        "    plt.figure(figsize=(15,4), dpi=dpi)\n",
        "    plt.subplot(1,3,1)\n",
        "\n",
        "  plt.title('Predicted lDDT')\n",
        "  plt.plot(O[\"plddt\"])\n",
        "  if Ls is not None:\n",
        "    L_prev = 0\n",
        "    for L_i in Ls[:-1]:\n",
        "      L = L_prev + L_i\n",
        "      L_prev += L_i\n",
        "      plt.plot([L,L],[0,100],color=\"black\")\n",
        "  plt.xlim(0,O[\"plddt\"].shape[0])\n",
        "  plt.ylim(0,100)\n",
        "  plt.ylabel('plDDT')\n",
        "  plt.xlabel('position')\n",
        "  plt.subplot(1,4 if \"lm_contacts\" in O else 3,2)\n",
        "\n",
        "  plt.title('Predicted Aligned Error')\n",
        "  Ln = O[\"pae\"].shape[0]\n",
        "  plt.imshow(O[\"pae\"],cmap=\"bwr\",vmin=0,vmax=30,extent=(0, Ln, Ln, 0))\n",
        "  if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "  plt.colorbar()\n",
        "  plt.xlabel('Scored residue')\n",
        "  plt.ylabel('Aligned residue')\n",
        "\n",
        "  if \"lm_contacts\" in O:\n",
        "    plt.subplot(1,4,3)\n",
        "    plt.title(\"contacts from LM\")\n",
        "    plt.imshow(O[\"lm_contacts\"],cmap=\"Greys\",vmin=0,vmax=1,extent=(0, Ln, Ln, 0))\n",
        "    if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "    plt.subplot(1,4,4)\n",
        "  else:\n",
        "    plt.subplot(1,3,3)\n",
        "  plt.title(\"contacts from Structure Module\")\n",
        "  plt.imshow(O[\"sm_contacts\"],cmap=\"Greys\",vmin=0,vmax=1,extent=(0, Ln, Ln, 0))\n",
        "  if Ls is not None and len(Ls) > 1: plot_ticks(Ls)\n",
        "  return plt\n",
        "\n",
        "plot_confidence(best_output, Ls=lengths, dpi=dpi)\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "#@title download predictions\n",
        "from google.colab import files\n",
        "os.system(f\"zip {ID}.zip {ID}/*\")\n",
        "files.download(f'{ID}.zip')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "AVn1z6H9o7GT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title animate outputs (optional)\n",
        "#@markdown if more than one sample generated\n",
        "import os\n",
        "if not os.path.isdir(\"colabfold\"):\n",
        "  os.system(\"git clone https://github.com/sokrypton/ColabFold.git\")\n",
        "import ColabFold.beta.colabfold as cf\n",
        "\n",
        "dpi =  100#@param {type:\"integer\"}\n",
        "color_by_plddt = True #@param {type:\"boolean\"}\n",
        "use_pca = True\n",
        "cycle = True \n",
        "\n",
        "import matplotlib\n",
        "from matplotlib import animation\n",
        "from IPython.display import HTML\n",
        "from sklearn.decomposition import PCA\n",
        "\n",
        "def mk_animation(xyz, labels, plddt, ref=0, Ls=None, line_w=2.0, dpi=100,color_by_plddt=False):\n",
        "\n",
        "  def ca_align_to_last(positions, ref):\n",
        "    def align(P, Q):\n",
        "      if Ls is None or len(Ls) == 1:\n",
        "        P_,Q_ = P,Q\n",
        "      else:\n",
        "        # align relative to first chain\n",
        "        P_,Q_ = P[:Ls[0]],Q[:Ls[0]]\n",
        "      p = P_ - P_.mean(0,keepdims=True)\n",
        "      q = Q_ - Q_.mean(0,keepdims=True)\n",
        "      return ((P - P_.mean(0,keepdims=True)) @ cf.kabsch(p,q)) + Q_.mean(0,keepdims=True)\n",
        "    \n",
        "    pos = positions[ref] - positions[ref].mean(0,keepdims=True)\n",
        "    best_2D_view = pos @ cf.kabsch(pos,pos,return_v=True)\n",
        "\n",
        "    new_positions = []\n",
        "    for i in range(len(positions)):\n",
        "      new_positions.append(align(positions[i],best_2D_view))\n",
        "    return np.asarray(new_positions)\n",
        "\n",
        "  # align to reference \n",
        "  pos = ca_align_to_last(xyz, ref)\n",
        "\n",
        "  fig, (ax1) = plt.subplots(1)\n",
        "  fig.set_figwidth(5)\n",
        "  fig.set_figheight(5)\n",
        "  fig.set_dpi(dpi)\n",
        "\n",
        "  xy_min = pos[...,:2].min() - 1\n",
        "  xy_max = pos[...,:2].max() + 1\n",
        "\n",
        "  for ax in [ax1]:\n",
        "    ax.set_xlim(xy_min, xy_max)\n",
        "    ax.set_ylim(xy_min, xy_max)\n",
        "    ax.axis(False)\n",
        "\n",
        "  ims=[]\n",
        "  for l,pos_,plddt_ in zip(labels,pos,plddt):\n",
        "    if color_by_plddt:\n",
        "      img = cf.plot_pseudo_3D(pos_, c=plddt_, cmin=50, cmax=90, line_w=line_w, ax=ax1)    \n",
        "    elif Ls is None or len(Ls) == 1:\n",
        "      img = cf.plot_pseudo_3D(pos_, ax=ax1, line_w=line_w)\n",
        "    else:\n",
        "      c = np.concatenate([[n]*L for n,L in enumerate(Ls)])\n",
        "      img = cf.plot_pseudo_3D(pos_, c=c, cmap=cf.pymol_cmap, cmin=0, cmax=39, line_w=line_w, ax=ax1)\n",
        "    ims.append([cf.add_text(f\"{l}\", ax1),img])\n",
        "    \n",
        "  ani = animation.ArtistAnimation(fig, ims, blit=True, interval=120)\n",
        "  plt.close()\n",
        "  return ani.to_html5_video()\n",
        "\n",
        "labels = np.array([f\"seed:{x}\" for x in range(len(traj))])\n",
        "pos = np.array([x[\"xyz\"] for x in traj])\n",
        "plddt = np.array([x[\"plddt\"] for x in traj])\n",
        "if use_pca and pos.shape[0] > 1:\n",
        "  pos_ca = pos\n",
        "  if lengths is not None and len(lengths) > 1:\n",
        "    pos_ca = pos_ca[:,:lengths[0]]\n",
        "  i,j = np.triu_indices(pos_ca.shape[1],k=1)\n",
        "  pos_ca_dm = np.sqrt(np.square(pos_ca[:,None,:,:] - pos_ca[:,:,None]).sum(-1))[:,i,j]\n",
        "  pc = PCA(1).fit_transform(pos_ca_dm)[:,0].argsort()\n",
        "  if cycle:\n",
        "    pc = np.append(pc,pc[1:-1][::-1])\n",
        "  pos = pos[pc]\n",
        "  labels = labels[pc]\n",
        "  plddt = plddt[pc]\n",
        "HTML(mk_animation(pos, labels, plddt, Ls=lengths, dpi=dpi, color_by_plddt=color_by_plddt))"
      ],
      "metadata": {
        "cellView": "form",
        "id": "NY8H_iM5WFUu"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "gpuClass": "standard",
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}